#include "motion_compensation_node.h"
#include "tic_toc.h"
#include "mathused.h"

EventProcessNode::EventProcessNode(ros::NodeHandle &nh)
	: nh_(nh), it_(nh)
{
	// rpg-mono
	K << 199.0923665423112, 0.0, 132.1920713777002, 0.0, 198.8288204700886, 110.7126600112956, 0.0, 0.0, 1.0;
	
	// FPV
	// K << 172.98992850734132, 0, 163.33639726024606, 0, 172.98303181090185, 134.99537889030861, 0, 0, 1;
	K_inverse = K.inverse();
	window_size_ = nh_.param("window_size_", 30000);
	median_depth_ = nh_.param("median_depth_", 0.5);
	map_blur_ = nh_.param("map_blur_", 5);
	events_sub_ = nh_.subscribe("/dvs/events", 2, &EventProcessNode::eventCallback, this);
	img_sub_ = nh_.subscribe("/dvs/image_raw", 2, &EventProcessNode::imageCallback, this);
	ba_sub_ = nh_.subscribe("/vins_estimator/bias_accl", 2, &EventProcessNode::baCallback, this);
	bg_sub_ = nh_.subscribe("/vins_estimator/bias_gyro", 2, &EventProcessNode::bgCallback, this);

	event_image_pub_ = it_.advertise("/mc/event_image", 1);
	raw_image_pub_ = it_.advertise("/mc/image_raw", 1);

	EventAccumulatedImage_ = Mat::zeros(sensor_rows, sensor_cols, CV_8U);
	events_buffer_.clear();
	imu_buffer_.clear();

	imu_sub_.subscribe(nh_, "/dvs/imu", 500, ros::TransportHints().tcpNoDelay());
	imu_cache_.reset(new message_filters::Cache<sensor_msgs::Imu>(imu_sub_, 100));
	imu_cache_->registerCallback(boost::bind(&EventProcessNode::imuCallback, this, _1));

	b_w_ << 0, 0, 0;
	b_a_ << 0, 0, 0;

	first_event = true;
	event_size_ = 0;
	ROS_INFO("Initialized!");
	ROS_INFO("window_size_: %d", window_size_);
	ROS_INFO("median_depth_: %f", median_depth_);
}

EventProcessNode::~EventProcessNode()
{
	event_image_pub_.shutdown();
	raw_image_pub_.shutdown();
}

void EventProcessNode::eventCallback(const dvs_msgs::EventArray::ConstPtr &emsg)
{
	// ROS_INFO("Callback entered!");
	// if (first_event)
	// {
	// 	t0 = emsg->events[0].ts;
	// 	first_event = false;
	// 	for (int i = 0; i < emsg->events.size(); i++)
	// 		events_buffer_.push_back(emsg->events[i]);
	// 	event_size_ = emsg->events.size();
	// 	return;
	// }
	// for (int i = 0; i < emsg->events.size(); i++)
	// {
	// 	events_buffer_.push_back(emsg->events[i]);
	// 	event_size_++;
	// 	// if (event_size_ >= window_size_)
	// 	// {
	// 	// 	t1 = emsg->events[i].ts;
	// 	// 	TicToc t_s;
	// 	// 	std::vector<sensor_msgs::Imu::ConstPtr> imu_vector = imu_cache_->getInterval(t0, t1);
	// 	// 	std::tie(R10, t10) = imuIntegrate(imu_vector);
	// 	// 	// ROS_INFO("R10: %f %f %f; t10: %f %f %f", R10[0], R10[1], R10[2], t10[0], t10[1], t10[2]);
	// 	// 	deltaT = (t1 - t0).toSec();
	// 	// 	// ROS_INFO("deltaT = %f", deltaT);
	// 	// 	// no_motion_compensate();
	// 	// 	motion_compensate();
	// 	// 	// ROS_INFO("t0 = %f t1 = %f", t0.toSec(), t1.toSec());
	// 	// 	double t = t_s.toc();
	// 	// 	// ROS_INFO("Precessing time: %lf ms; Event Size: %d; Average rate: %lf event/ms", t, event_size_, event_size_ / t);
	// 	// 	events_buffer_.clear();
	// 	// 	t0 = t1;
	// 	// 	event_size_ = 0;
	// 	// }
	// }
	for (const auto &e : emsg->events)
		events_buffer_.push_back(e);
	event_size_ += emsg->events.size();
}

void EventProcessNode::imageCallback(const sensor_msgs::ImageConstPtr &img_msg)
{
	image_time = img_msg->header.stamp;
	if (event_size_ >= window_size_)
	{
		// ROS_INFO("event_size_ = %d", event_size_);
		t1 = events_buffer_[event_size_ - 1].ts;
		// t1 = image_time;
		int res = event_size_ - window_size_;
		t0 = events_buffer_[res].ts;
		std::vector<sensor_msgs::Imu::ConstPtr> imu_vector = imu_cache_->getInterval(t0, t1);
		std::tie(R10, t10) = imuIntegrate(imu_vector);
		deltaT = (t1 - t0).toSec();
		motion_compensate();
		events_buffer_.erase(events_buffer_.begin(), events_buffer_.begin() + res);
		event_size_ = window_size_;
	}
	else
		ROS_INFO("skip event_frame");

	cv_bridge::CvImageConstPtr ptr;
	ptr = cv_bridge::toCvCopy(img_msg, sensor_msgs::image_encodings::MONO8);
	raw_image_pub_.publish(ptr->toImageMsg());
}

void EventProcessNode::imuCallback(const sensor_msgs::Imu::ConstPtr &imu_msg)
{
}

void EventProcessNode::baCallback(const motion_compensation::biasPtr &msg)
{
	b_a_[0] = (*msg).x;
	b_a_[1] = (*msg).y;
	b_a_[2] = (*msg).z;
}

void EventProcessNode::bgCallback(const motion_compensation::biasPtr &msg)
{
	b_w_[0] = (*msg).x;
	b_w_[1] = (*msg).y;
	b_w_[2] = (*msg).z;
}

void EventProcessNode::motion_compensate()
{
	std::lock_guard<std::mutex> lock_events(data_mutex_);
	// for (int i = 0; i < event_size_; i++)
	int i = event_size_ - window_size_;
	for (; i < event_size_; i++)
	{
		// ROS_INFO("i= %d", i);
		auto e = events_buffer_[i];
		tt = e.ts;
		// ROS_INFO("tt = %f t0 = %f t1 = %f", tt.toSec(), t0.toSec(), t1.toSec());
		if (tt <= t0 || tt > t1)
			continue;

		double delta_t_ = (tt - t0).toSec();
		re = (delta_t_ / deltaT) * R10;
		te = (delta_t_ / deltaT) * t10;
		rot = exp_SO3(re);

		eventVec[0] = e.x;
		eventVec[1] = e.y;
		eventVec[2] = 1;
		// ROS_INFO("Before motion compensation: e.x = %d, e.y = %d", e.x, e.y);
		eventVec = K_inverse * eventVec;
		eventVec *= median_depth_;
		eventVec = rot * eventVec + te;
		eventVec[0] /= eventVec[2];
		eventVec[1] /= eventVec[2];
		eventVec[2] = 1;
		eventVec = K * eventVec;

		int ix = static_cast<int>(eventVec[0]);
		int iy = static_cast<int>(eventVec[1]);

		// ROS_INFO("After motion compensation: ix = %d, iy = %d", ix, iy);
		if (IsWithinTheBoundary(ix, iy))
			EventAccumulatedImage_.at<uchar>(static_cast<int>(iy), static_cast<int>(ix)) += 1;
			// EventAccumulatedImage_.at<uchar>(static_cast<int>(iy), static_cast<int>(ix)) = 1;
	}

	cv::normalize(EventAccumulatedImage_, EventAccumulatedImage_, 0, 255, cv::NORM_MINMAX);
	// cv::GaussianBlur(EventAccumulatedImage_, EventAccumulatedImage_, cv::Size(map_blur_, map_blur_), 0.);
	static cv_bridge::CvImage cv_image;
	cv_image.encoding = "mono8";
	cv_image.image = EventAccumulatedImage_.clone();
	// cv_image.header.stamp = t1;
	cv_image.header.stamp = image_time;
	event_image_pub_.publish(cv_image.toImageMsg());
	EventAccumulatedImage_.setTo(0);
}

void EventProcessNode::no_motion_compensate()
{
	for (int i = 0; i < event_size_; i++)
	{
		/* 不补偿 */
		auto e = events_buffer_[i];
		// ROS_INFO("tt = %lf t0 = %lf t1 = %lf", e.ts.toSec(), t0.toSec(), t1.toSec());
		if (IsWithinTheBoundary(e.x, e.y))
			EventAccumulatedImage_.at<uchar>(static_cast<int>(e.y), static_cast<int>(e.x)) += 1;
	}
	// ROS_INFO("Motion Compensate entered!");
	cv::normalize(EventAccumulatedImage_, EventAccumulatedImage_, 0, 255, cv::NORM_MINMAX);
	// cv::Mat undistorted_events_ = undistortImage(EventAccumulatedImage_);
	static cv_bridge::CvImage cv_image;
	cv_image.encoding = "mono8";
	cv_image.image = EventAccumulatedImage_.clone();
	cv_image.header.stamp = t1;
	event_image_pub_.publish(cv_image.toImageMsg());
	EventAccumulatedImage_.setTo(0);
}

std::tuple<Eigen::Vector3f, Eigen::Vector3f> EventProcessNode::imuIntegrate(std::vector<sensor_msgs::Imu::ConstPtr> &imu_vector) const
{
	// 一段时间里的imu观测
	int num_measurement = imu_vector.size();
	Eigen::Matrix3f R_meas;
	Eigen::Vector3f v_meas;
	Eigen::Vector3f p_meas;
	R_meas << 1, 0, 0, 0, 1, 0, 0, 0, 1;
	v_meas << 0, 0, 0;
	p_meas << 0, 0, 0;
	// 获取imu测量
	for (int i = 0; i < num_measurement - 1; i++)
	{
		// 获取陀螺仪和加速度计测量值
		Eigen::Vector3f w;
		Eigen::Vector3f a;
		w(0) = imu_vector[i]->angular_velocity.x;
		w(1) = imu_vector[i]->angular_velocity.y;
		w(2) = imu_vector[i]->angular_velocity.z;
		a(0) = imu_vector[i]->linear_acceleration.x;
		a(1) = imu_vector[i]->linear_acceleration.y;
		a(2) = imu_vector[i]->linear_acceleration.z;
		// 去除偏差(预设值,不参与更新)
		w = w - b_w_;
		a = a - b_a_;
		float delta_t = (imu_vector[i + 1]->header.stamp - imu_vector[i]->header.stamp).toSec();
		Eigen::Matrix<float, 3, 3> delta_R_jm1_j = exp_SO3<float>(w * delta_t);
		// 更新状态量
		R_meas *= delta_R_jm1_j;
		v_meas += R_meas * a * delta_t;
		p_meas += v_meas * delta_t + 0.5 * R_meas * a * delta_t * delta_t;
	}
	Eigen::Vector3f R_ret = log_SO3<float>(R_meas);
	// 返回C++元组
	return std::make_tuple(R_ret, p_meas);
}

inline bool EventProcessNode::IsWithinTheBoundary(const int &x, const int &y)
{
	return (x >= 0 && x < sensor_cols && y >= 0 && y < sensor_rows);
}

void EventProcessNode::unproject(Eigen::Vector3f &vec)
{
	vec = K_inverse * vec;
	vec *= median_depth_;
}

void EventProcessNode::project(Eigen::Vector3f &vec)
{
	vec[0] /= vec[2];
	vec[1] /= vec[2];
	vec[2] = 1;
	vec = K * vec;
}

int main(int argc, char **argv)
{
	ros::init(argc, argv, "motion_compensation_node");
	ros::NodeHandle nh("~");
	EventProcessNode event_process_node(nh);
	ros::spin();

	return 0;
}
